#include "IAP.h"
#include "SEGGER_RTT.h"
#include "main.h"
#include "tim.h"
#include "usart.h"

typedef void (*jump_func)(void);

// clang-format off

/**
 * @brief show address and size information of bootloader and application by RTT.
 */
void show_boot_info(void)
{ 
    SEGGER_RTT_printf(0, "---------- Enter BootLoader ----------\r\n");
    SEGGER_RTT_printf(0, "\r\n");
    SEGGER_RTT_printf(0, "======== flash pration table =========\r\n");
    SEGGER_RTT_printf(0, "| name     | offset     | size       |\r\n");
    SEGGER_RTT_printf(0, "--------------------------------------\r\n");
    SEGGER_RTT_printf(0, "| boot     | 0x%08X | 0x%08X |\r\n", BOOT_START_ADDR, BOOT_FLASH_SIZE);
    SEGGER_RTT_printf(0, "| app      | 0x%08X | 0x%08X |\r\n", APP_START_ADDR, APP_FLASH_SIZE);
    SEGGER_RTT_printf(0, "======================================\r\n");
}

// clang-format on

/**
 * @brief jump to application and run
 * @param[in] app_addr application address
 * @return result.
 */
uint8_t jump_app(uint32_t app_addr)
{
    uint32_t app_msp = *(volatile uint32_t*)(app_addr);
    uint32_t jump_addr = *(volatile uint32_t*)(app_addr + 4);
    if ((app_msp & 0x2FFE0000) == 0x20000000) {
        jump_func JumpToApp = (jump_func)jump_addr;
        __set_MSP(app_msp);
        JumpToApp();

        return 1;
    }
    return 0;
}

/**
 * @brief get flash_sector by flash_address
 * @param[in] addr flash address
 * @return flash sector
 */
uint16_t FLASH_GetSector(uint32_t addr)
{
    if (addr < ADDR_FLASH_SECTOR_1)
        return FLASH_SECTOR_0;
    else if (addr < ADDR_FLASH_SECTOR_2)
        return FLASH_SECTOR_1;
    else if (addr < ADDR_FLASH_SECTOR_3)
        return FLASH_SECTOR_2;
    else if (addr < ADDR_FLASH_SECTOR_4)
        return FLASH_SECTOR_3;
    else if (addr < ADDR_FLASH_SECTOR_5)
        return FLASH_SECTOR_4;
    else if (addr < ADDR_FLASH_SECTOR_6)
        return FLASH_SECTOR_5;
    else if (addr < ADDR_FLASH_SECTOR_7)
        return FLASH_SECTOR_6;
    else
        return FLASH_SECTOR_7;
    return 0xff;
}

/* ========================================================================================= */
/*                                           IAP逻辑                                          */
/* ========================================================================================= */

uint8_t app_buff[20 * 1024]; // 20K
uint16_t app_len = 0;        // app length

// 用于软复位重启的标志，该变量不会被零初始化
// 升级成功后，写入0xff112233，然后软复位
uint32_t jump_init __attribute__((used, section(".bss.NoInit")));

uint32_t update_app()
{
    uint16_t ret = 0;
    uint8_t key_press = 0;

    show_boot_info();
    __HAL_TIM_ENABLE(&htim10);
    __HAL_TIM_CLEAR_IT(&htim10, TIM_IT_UPDATE);
    // if press key in 5s, update app
    // else, jump to app directly.
    while (!__HAL_TIM_GET_FLAG(&htim10, TIM_FLAG_UPDATE)) {
        if (HAL_GPIO_ReadPin(KEY_GPIO_Port, KEY_Pin) == GPIO_PIN_RESET) {
            HAL_Delay(20);
            if (HAL_GPIO_ReadPin(KEY_GPIO_Port, KEY_Pin) == GPIO_PIN_RESET) {
                key_press = 1;
                break;
            }
        }
    }

    if (key_press) {
        // receive app bin
        SEGGER_RTT_printf(0, "Please send your bin file by UART\n");
        ret = HAL_UARTEx_ReceiveToIdle(&huart1, app_buff, 20 * 1024, &app_len, HAL_MAX_DELAY);

        // if receive ok
        if (ret == HAL_OK) {
            // erase flash
            FLASH_EraseInitTypeDef flash_erase = {0};
            flash_erase.VoltageRange = FLASH_VOLTAGE_RANGE_3;
            flash_erase.Banks = FLASH_BANK_1;
            flash_erase.Sector = FLASH_GetSector(APP_START_ADDR);
            flash_erase.NbSectors = FLASH_GetSector(APP_START_ADDR + app_len - 1) -
                                    FLASH_GetSector(APP_START_ADDR) + 1;
            flash_erase.TypeErase = FLASH_TYPEERASE_SECTORS;

            uint32_t sector_error = 0;
            HAL_FLASH_Unlock();
            ret = HAL_FLASHEx_Erase(&flash_erase, &sector_error);
            HAL_FLASH_Lock();
            SEGGER_RTT_printf(0, "erase app pration !!!\n");

            // if erase successfully
            if ((ret == HAL_OK) && (sector_error == 0xFFFFFFFF)) {
                // write app bin to flash
                HAL_FLASH_Unlock();
                uint32_t address = APP_START_ADDR;
                for (uint32_t i = 0; i < app_len; i++) {
                    HAL_FLASH_Program(FLASH_TYPEPROGRAM_BYTE, address, app_buff[i]);
                    address++;
                }
                HAL_FLASH_Lock();
                SEGGER_RTT_printf(0, "write %d bytes to app pration!!!\n", app_len);
                HAL_Delay(100);

                return 0;
            }
        }
    }

    return 1;
}

void iap_function()
{
    uint16_t ret = 0xff;

    if (jump_init == 0xff112233) {
        SEGGER_RTT_printf(0, "Var jump_init is %x!!!\n", jump_init);
        HAL_Delay(100);
        jump_init = 0x00000000;
        ret = jump_app(APP_START_ADDR);
        if (ret == 0) {
            SEGGER_RTT_printf(0, "jump to application failed!!!\n");
        }
    }
    else {
        SEGGER_RTT_printf(0, "Var jump_init is NoInit!!!\n");
        ret = update_app();
        if (ret == 0) { // update app successfully
            jump_init = 0xff112233;
            HAL_NVIC_SystemReset();
        }
        else { // receive app bin timeout or failed
            SEGGER_RTT_printf(0, "Update timeout or failed!!!\n", jump_init);
            HAL_Delay(100);
            jump_app(APP_START_ADDR);
        }
    }
}
